# -*- coding: utf-8 -*-
"""
Created on Thu May  8 12:23:48 2014

@author: dorme

The basic idea is to take an Excel file containing a list of events
and output an ICS file of the events along with a simple html document
showing a weekly table and description list that can be put on Blackboard
or imported into Excel or Word
"""

import icalendar
import openpyxl
import datetime as dt
import pytz
import os, sys
import types
import argparse  

class Event:

    """
    Define a class that takes an iter imported row from the calender Excel
    format and adds some processing steps and a simple column span method
    """

    def __init__(self, evRow):

        """
        the file contains a row of cells in this order, which are
        imported as a tuple of RawCell objects in this order:
        Term, TermWeek, YearWeek, WeekDay, Date, Day, Start, End,
        Location, Year, Group, Module, SessionCode, Description, Staff
        """

        # get row for error reporting
        rowNum = evRow[0].row

        # sanity testing
        if len(evRow) != 14:
            sys.stderr.write('Number of columns in Timetable sheet not ' +
                             'equal to 15\nTry clearing cells to right of ' +
                             'timetable\n')
            sys.exit() # TODO think about exceptions here

        # extract the data and types
        evCols = [x.column for x in evRow]
        evRow = [x.internal_value for x in evRow]
        evTypes = [type(x) for x in evRow]

        # test the types - None allows blank cells
        validTypes = [[unicode, str],  # Term
                      [float, int],  # Academic Week
                      [float, int],  # Week Sequence
                      [float, int],  # Day number
                      [dt.datetime],  # Date
                      [unicode, str],  # Day
                      [dt.time],  # Start
                      [dt.time],  # End
                      [unicode, str, float, int, types.NoneType],  # Location
                      [unicode, str, types.NoneType],  # Module
                      [unicode, str, types.NoneType],  # Session Code
                      [unicode, str, types.NoneType],  # Group
                      [unicode, str],  # Description
                      [unicode, str, types.NoneType]]  # Staff

        matchTypes = [i in v for i, v in zip(evTypes, validTypes)]

        if not all(matchTypes):
            badCols = [c for c, b in zip(evCols, matchTypes) if not b]
            sys.stderr.write('Value type problem in row ' + str(rowNum) +
                             ', columns ' + ','.join(badCols) + '.\n' +
                             'Check strings, numbers and dates are as ' +
                             'expected or not blank.\n')
            print evRow
            sys.exit()
            # TODO think about exceptions here to give full report
            # rather than error by error

        # screen out unicode
        evRow = [x.encode('ascii', 'xmlcharrefreplace') if type(x)
                 is unicode else x for x in evRow]

        # check types
        evTypes = [type(x) for x in evRow]

        # strip off trailing space
        evRow = [x.strip() if type(x) is str else x for x in evRow]

        # name the basic columns and enforce types
        self.Term = evRow[0]
        self.TermWeek = int(evRow[1])
        self.YearWeek = int(evRow[2])
        self.WeekDay = int(evRow[3])
        self.Date = evRow[4]
        self.Day = evRow[5]
        self.Start = evRow[6]
        self.End = evRow[7]
        # handle string, number and missing for location
        if evRow[8] is None:
            self.Location = None
        elif type(evRow[8]) in (int, float):
            self.Location = str(int(evRow[8]))
        else:
            self.Location = evRow[8]
        self.Module = evRow[9]
        self.SessionCode = evRow[10]
        self.Group = evRow[11]
        self.Description = evRow[12]
        self.Staff = evRow[13]

        # event title - module code + group or just description if missing
        if self.Module is None:
            self.Title = self.Description
            self.Description = ''
        else:
            self.Title = '%s.%s' % (self.Module, self.SessionCode)
        if self.Group is not None:
            self.Title = ' [' + self.Group + '] ' + self.Title

        # add bracketed info
        glsLab = ['', 'Group: ', 'Staff: ']
        glsVal = [self.Location, self.Group, self.Staff]
        gls = zip(glsLab, glsVal)
        gls = [i + v for i, v in gls if v is not None]
        if gls != []:
            gls = '  [' + '; '.join(gls) + ']'
            self.Description += gls

        # time range as string
        self.TimeRange = '%s &ndash; %s' % (self.Start.strftime('%H:%M'),
                                            self.End.strftime('%H:%M'))
        # initalise column span, RGB, TZ using defaults
        self.colSpan()
        self.colourLookup()
        self.tzLocalize()

    def colSpan(self, colsPerHour=4, calStart=9, calEnd=17):

        # adds or updates the column span and start/end for entries
        # in the html version of the calendars.
        # col numbers start at 1, assuming a left hand row label in [0]

        # start and end time as decimal hours,
        # - trimming at low and high bounds
        evStart = float(self.Start.hour) + float(self.Start.minute)/60
        evStart = max(calStart, evStart)
        evEnd = float(self.End.hour) + float(self.End.minute)/60
        evEnd = min(evEnd, calEnd)

        self.FirstCol = int((evStart - calStart) // (1.0/colsPerHour)) + 1
        self.LastCol = int((evEnd - calStart) // (1.0/colsPerHour)) + 1
        self.Span = self.LastCol - self.FirstCol


    def colourLookup(self, colDict={'default': '#CCC'}):

        # color code - excel only handles #FFFFFF on html import
        if self.Module is None or self.Module not in colDict.keys():
            self.RGB = colDict['default']
        else:
            self.RGB = colDict[self.Module]

    def tzLocalize(self, tz=pytz.timezone('Europe/London')):

        # need time zone aware date+time for iCal format
        self.startDT = tz.localize(dt.datetime.combine(self.Date, self.Start))
        self.endDT = tz.localize(dt.datetime.combine(self.Date, self.End))


class tableStructure:

    def __init__(self, colsPerHour=4, calStart=9, calEnd=17):
        # gets a list of column start times and column indices
        try:
            self.startTime = dt.time(calStart)
        except:
            sys.stderr.write('Calendar start hour not an integer in 0..23\n')
            sys.exit()

        try:
            self.endTime = dt.time(calEnd)
        except:
            sys.stderr.write('Calendar end hour not an integer in 0..23\n')
            sys.exit()

        if self.startTime > self.endTime:
            sys.stderr.write('Calendar end hour earlier than start hour\n')
            sys.exit()

        self.colsPerHour = int(colsPerHour)
        self.colTimeSpan = dt.timedelta(minutes=60/float(self.colsPerHour))

        # convert to arbitrary datetime for timedelta methods
        baseDate = dt.datetime(1901, 1, 1)
        start = dt.datetime.combine(baseDate, self.startTime)
        end = dt.datetime.combine(baseDate, self.endTime)
        colTimes = ['']  # row label column on left

        while start < end:
            colTimes.append(start.time())
            start += self.colTimeSpan

        colNumbers = range(len(colTimes))

        self.colTimes = colTimes
        self.colNumbers = colNumbers
        self.nCols = len(self.colNumbers)

        self.hourLabels = [self.colTimes[i].strftime('%H:%M') for i in
                           range(1, self.nCols, self.colsPerHour)]

        self.setHeadWidth()

    def setHeadWidth(self, totalWidth=1000):

        # add a header and footer
        colStatements = '<COL width="%i">\n' % int(totalWidth/self.nCols)
        self.HeaderBorder = '<TABLE  cellpadding="2" style="border-collapse:' + \
                      ' collapse;border: 1px solid black;">\n' + \
                      colStatements * self.nCols
        self.HeaderNoBorder = '<TABLE  cellpadding="2" style="border-collapse:' + \
                      ' collapse;">\n' + \
                      colStatements * self.nCols
        self.Footer = '</TABLE>\n'


def splitEventList(eventList, attr):

    """
    Split a list of events into a dictionary of sub lists
    on one of the attributes
    """

    attrList = [getattr(ev, attr) for ev in eventList]
    attrSet = set(attrList)
    result_lists = dict((key, []) for key in attrSet)

    for i, v in zip(attrList, eventList):
        result_lists[i].append(v)

    return result_lists


def loadEvents(fname):

    # open the excel file to iterate over rows
    try:
        wb = openpyxl.load_workbook(fname, use_iterators=True, data_only=True)
    except openpyxl.shared.exc.InvalidFileException:
        sys.stderr.write('Excel file %s not found\n' % os.path.basename(fname))
        sys.exit()

    # get the sheets available
    nm = wb.get_sheet_names()

    # find and access the timetable sheet
    if 'Timetable' not in nm:
        sys.stderr.write("No worksheet called 'Timetable' found.\n")
        sys.exit()
    sh = wb.get_sheet_by_name('Timetable')

    # set up the row iterator
    rowIt = sh.iter_rows()

    # skip first two lines and check headers
    rowIt.next()
    rowIt.next()
    headers = rowIt.next()
    headers = [x.internal_value for x in headers]

    expected = [u'Term', u'Academic Week', u'Week Sequence', u'Day Number', 
                u'Date', u'Day', u'Start', u'End', u'Location',
                 u'Module', u'Session Code', u'Group', u'Details', u'Staff']
    if expected != headers:
        sys.stderr.write("Timetable headers in row 3 do " +
                         "not match expected headers.\n")
        sys.exit()

    # load the rest of the data into a list
    eventList = []
    for row in rowIt:
        eventList.append(Event(row))

    # look for a colour sheet
    colDict = {}
    foundModules = set([x.Module for x in eventList])

    if 'Module Colours' not in nm:
        sys.stderr.write("No 'Module Colours' worksheet found, " +
                         "using default.\n")
        colDict = {'default': '#AAA'}
    else:
        # load the module names and colours
        # - excel needs hex encoded colours
        colSheet = wb['Module Colours']

        rowIt = colSheet.iter_rows()
        row = rowIt.next()
        for row in rowIt:
            rgbTrip = row[1].internal_value.split(',')
            rgbTrip = [int(x) for x in rgbTrip]
            colDict[row[0].internal_value] = '#%02X%02X%02X' % tuple(rgbTrip)

        # check the colours are complete
        colModules = set(colDict.keys())

        if not colModules.issuperset(foundModules):
            sys.stderr.write("Modules in timetable missing from module" +
                             " colour sheet, using defaults:\n")
            for x in foundModules.difference(colModules):
                sys.stderr.write(x + '\n')

        if 'default' not in colDict.keys():
            colDict['default'] = '#CCC'

    [ev.colourLookup(colDict) for ev in eventList]
    return eventList


def exportICS(eventList, fname, calname, caldesc):

    # initialise the calendar and provide top level detail
    cal = icalendar.Calendar()

    cal.add('method', 'PUBLISH')
    cal.add('calscale', 'GREGORIAN')
    cal.add('x-wr-calname', calname)
    cal.add('x-wr-caldesc', caldesc)

    # populate the ICS file
    for ev in eventList:

        # create an event and add information
        # - don't use the attendee field and organiser role for staff
        #   because it triggers the scheduling mechanism
        event = icalendar.Event()

        # start and end date/time
        event.add('dtstart', ev.startDT)
        event.add('dtend', ev.endDT)

        # summary - module code + group
        event.add('summary', ev.Title)

        # description - description + staff
        event.add('description', ev.Description)
        event.add('location', ev.Location)

        # TODO (eventually) setting event colour not yet implemented in ICS
        # but is in some draft proposals

        # add that event to the calendar
        cal.add_component(event)

    # export the file
    # this could return the cal object
    f = open(fname, 'wb')
    f.write(cal.to_ical())
    f.close()


def htmlWeek(thisWeek, tabStruct):

    '''
    takes a list of events in a week and turns it into a list of:
    [[html for calendar rows in week], [html for details]]
    '''

    # create a week header row
    # - need to find the week range from the available days
    #   which might not include all days in the week
    refDay = thisWeek[0].Date
    refDayNum = refDay.weekday()
    monday = refDay - dt.timedelta(days=refDayNum)
    friday = refDay + dt.timedelta(days=4 - refDayNum)

    # get the events by day
    eventsByDay = splitEventList(thisWeek, 'Date')

    # 1)  get the details for each event as a table
    # templates
    dayRow = '<TD colspan="%i" style="text-align:left;">%s</TD>\n'
    eventRow = '<TD></TD>\n' + \
               '<TD colspan="%i" style="vertical-align:top;' + \
               'font-style:italic;">%s</TD>\n' + \
               '<TD colspan="%i"><B>%s</B>&nbsp;&nbsp;%s</TD>\n'

    # build up a list of rows
    weekDetails = []
    dayKeys = eventsByDay.keys()
    dayKeys.sort()
    for k in dayKeys:
        # add the day row at the top
        kFormat = k.strftime('%A %d %B %Y')
        weekDetails.append(dayRow % (tabStruct.nCols, kFormat))
        # add a row for each event
        for ev in eventsByDay[k]:
            timeSpan = int(tabStruct.colsPerHour * 1.5)
            descSpan = tabStruct.nCols - (timeSpan + 1)
            thisRow = eventRow % (timeSpan, ev.TimeRange,
                                  descSpan, ev.Title, ev.Description)
            weekDetails.append(thisRow)

    # 2) Pack the titles of each event into rows in a mini calendar

    # hard coded styling: want flexibility? Too bad!
    weekHdr = '<TD colspan="%i" style="text-align:right;' + \
              'border:solid 1px; background-color:#CCC;"><span' + \
              ' style="float:left">&nbsp;&nbsp;<B>Week %i</B></span>' + \
              '<i>%s &ndash; %s</i>&nbsp;&nbsp;</TD>\n'
    weekHdr = [weekHdr % (tabStruct.nCols, thisWeek[0].TermWeek,
                          monday.strftime('%d/%m/%Y'),
                          friday.strftime('%d/%m/%Y'))]

    # create a row of hour labels, including the left hand gap
    hourLabels = '<TD colspan="%i" style="text-align:left;' + \
                 'border:solid 1px;font-size:small;' + \
                 'background-color:#CCC;">%s</TD>\n'

    hourDat = zip([1] + [tabStruct.colsPerHour] * len(tabStruct.hourLabels),
                  [''] + tabStruct.hourLabels)
    hourLabels = [hourLabels % (c, t) for c, t in hourDat]

    # fill in any missing days in this week as empty lists
    daysInDict = eventsByDay.keys()
    thisWeekDays = [monday + dt.timedelta(days=x) for x in range(5)]
    missingDays = set(thisWeekDays).difference(set(daysInDict))
    for d in missingDays:
        eventsByDay[d] = []

    # pack the events in each day into rows in day order
    weekRows = [weekHdr, hourLabels]
    for k in thisWeekDays:
        rowsToAdd = packRows(eventsByDay[k], k, tabStruct)
        for rw in rowsToAdd:
            weekRows.append(rw)

    return {'cal': weekRows, 'details': weekDetails}


def packRows(dayList, day, tabStruct):

    # need to pass day because dayList might be empty
    dayList.sort(key=lambda x: x.Start)

    # takes a set of events for a day and packs them into a set of rows
    # as tuple of (text, colour, start row, span).
    rowEmptyFrom = [1]
    dayTuple = ('#CCC', day.strftime('%a'))
    dayCell = '<TD colspan="1" style="text-align:left;font-size:small;' + \
              'border:solid 1px;background-color:%s;">%s</TD>\n'
    eventCell = '<TD colspan="%i" style="text-align:center;' + \
                'border:solid 1px;background-color:%s;">%s</TD>\n'
    emptyCell = '<TD></TD>\n'

    rowContents = [[dayCell % dayTuple]]

    for ev in dayList:

        # check which row this will fit in
        willFit = [x <= ev.FirstCol for x in rowEmptyFrom]
        firstFit = willFit.index(True) if True in willFit else None

        # add a new row if there's no row it will fit in
        if firstFit is None:
            rowEmptyFrom.append(1)
            rowContents.append([dayCell % dayTuple])
            firstFit = len(rowContents) - 1

        # pack out the appropiate row with empty cells if needed
        gapLength = ev.FirstCol - rowEmptyFrom[firstFit]
        if gapLength > 0:
            rowContents[firstFit].append(emptyCell * gapLength)
        # insert the event into the relevant row
        evTuple = (ev.Span, ev.RGB, ev.Title)
        rowContents[firstFit].append(eventCell % evTuple)
        rowEmptyFrom[firstFit] = ev.LastCol

    # finish packing row ends
    for i, lastCell in enumerate(rowEmptyFrom):
        gapLength = tabStruct.nCols - lastCell
        if gapLength > 0:
            rowContents[i].append(emptyCell * gapLength)

    return rowContents


def htmlOutWeek(conn, weekHTML):
    # simple function to concatenate all the text in a week block
    # of html and dump it out to a file connection
    rows = [''.join(x) for x in weekHTML]
    rows = ['<TR>\n%s</TR>\n' % x for x in rows]
    conn.write(''.join(rows))


def main(argv):

    # set up the parser and show defaults
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    
    # some input checking functions
    def IWG(value):
        
        value = str(value).upper()
        value = set(list(value))
        valid = value.issubset(set(['W','I','G']))

        if not valid:
             raise argparse.ArgumentTypeError("Output format option can " +
                                              "only contain I, W and G")
        return value
                
    # positional argument inputs
    parser.add_argument('file', nargs='+', help="Path to Excel file(s)" +
                        ' containig timetable details') 
    
    # optional keyword arguments
    parser.add_argument("-r", "--output-root", type=str,
                        dest='outputRoot', default='',
                        help='Filename root for output files.')
    parser.add_argument("-o", "--output-format", type=IWG, default='IWG', 
                        help='String of output types. I: .ics calendar ' +
                             'file, G: HTML grid calendar, W: HTML ' +
                             'weekly grid and details.')
    parser.add_argument("-m", "--modules", type=str, default='',
                        help='Comma separated list giving a subset ' +
                             'of modules to extract from files. The default ' +
                             '("") is to include all events')
    parser.add_argument("-x", "--explode", action='store_true', 
                        dest='explode', default=False,
                        help='Export modules to separate files.')
    parser.add_argument("-c", "--cols-per--hour", type=int, 
                        dest='colsPerHour', default=4,
                        help='Resolution of grid calendar.')
    parser.add_argument("-s", "--start", type=int,
                        dest='calStart', default=9,
                        help='Earliest hour in grid calendar.')
    parser.add_argument("-e", "--end", type=int,
                        dest='calEnd', default=17,
                        help='Latest hour in grid calendar.')

    args = parser.parse_args()

    # get a  timestamp and source file list to include in files
    now = dt.datetime.today()
    timestamp = 'Timestamp: ' + now.strftime('%H:%M %Y-%m-%d') + '\n'
    fileBase = [os.path.basename(x) for x in args.file]
    sourceFiles = 'Source file(s): ' + ', '.join(fileBase) + '\n'

    # get a default file root - first file base
    if args.outputRoot == '':
        rootname = os.path.splitext(os.path.basename(args.file[0]))[0]
        if len(args.file) > 1:
            rootname = rootname + "AndOthers"
    else:
        rootname = args.outputRoot
    
    # load the event table from all named files
    eventList = []
    for f in args.file:
        thisFile = loadEvents(f)
        sys.stdout.write('Imported %i events from %s\n' 
                          % (len(thisFile), os.path.basename(f)))
        eventList += thisFile

    # subset to the modules requested?
    if args.modules is not '':
        foundModules = set([x.Module for x in eventList])
        requestedModules = set(args.modules.split(','))
        if requestedModules.issubset(foundModules):
            eventList = [ev for ev in eventList 
                         if ev.Module in requestedModules]
        else:
            sys.stderr.write('Requested modules (%s) not all present ' +
                             'in file(s).\n' % args.modules)
            sys.exit()

    # get the table column structure and update the events 
    tabStruct = tableStructure(args.colsPerHour, args.calStart, args.calEnd)
    for ev in eventList:
        ev.colSpan(args.colsPerHour, args.calStart, args.calEnd)

    # Are we exploding the events into a dictionary by module
    if args.explode:
        eventList = splitEventList(eventList, 'Module')
        # handle None for events mising a module name
        if None in eventList.keys():
            eventList['NoModule'] = eventList.pop(None)
        suffix = ['_' + x for x in eventList.keys()] 
    else:
        eventList = {'all': eventList}  # Stuffed in a dictionary to match
        suffix = ['']
 
    # loop the event list
    for i, k in enumerate(eventList.keys()):
        
        # set up file names
        icsFilename = rootname + suffix[i] + ".ics"
        gridFilename = rootname + suffix[i] + "_Grid.html"
        weekFilename = rootname + suffix[i] + "_Week.html"
    
        # do we want the ICS file
        if 'I' in args.output_format:
            exportICS(eventList[k], icsFilename, rootname + '_' + k, 
                      sourceFiles + timestamp)
    
        # do we want html?
        if 'W' in  args.output_format or 'G' in args.output_format:
            weekList = splitEventList(eventList[k], 'TermWeek')
            weekKeys = weekList.keys()
            weekKeys.sort()
            weekHTMLRows = [htmlWeek(weekList[k], tabStruct) for k in weekKeys]

        # do we export the grid table
        if 'G' in args.output_format:
            of = open(gridFilename, 'w')
            of.write('<!DOCTYPE html>\n<html>\n<head>\n<title>')
            of.write(rootname + suffix[i] +'</title>\n</head>\n<body>\n')
            of.write(sourceFiles + '<br>\n' + timestamp + '<br>\n')
            of.write(tabStruct.HeaderBorder)
            for week in weekHTMLRows:
                htmlOutWeek(of, week['cal'])
            of.write(tabStruct.Footer)
            of.write('</body>\n</html>\n')
            of.close()
    
        # do we export the full timetable
        if 'W' in args.output_format:
            of = open(weekFilename, 'w')
            of.write('<!DOCTYPE html>\n<html>\n<head>\n<title>')
            of.write(rootname + suffix[i] +'</title>\n</head>\n<body>\n')
            of.write(sourceFiles + '<br>\n' + timestamp + '<br><br>\n')
            for week in weekHTMLRows:
                of.write(tabStruct.HeaderBorder)
                htmlOutWeek(of, week['cal'])
                of.write(tabStruct.Footer)
                of.write('<BR>\n')
                of.write(tabStruct.HeaderNoBorder)
                htmlOutWeek(of, week['details'])
                of.write(tabStruct.Footer)
                of.write('<BR>\n')
            of.write('</body>\n</html>\n')
            of.close()

if __name__ == '__main__':
    sys.exit(main(sys.argv))
